/*
 * KiDebugPort.cpp
 *
 *  Created on: 2013-10-25
 *      Author: WinDDK
 */

#include "KiDebug.h"

#define MAKELONG(a, b) ((unsigned long) (((unsigned short) (a)) | ((unsigned long) ((unsigned short) (b))) << 16))

typedef struct _CODE_SIGN
{
    PCHAR Signture;
    ULONG Offset;
    PCHAR Description;

} CODE_SIGN;

CODE_SIGN DebugPortSign[] =
{
{ "64??????????8B????39B8????????74??", 9 + 2, "KiDispatchException" },
{ "8B??????????8B81????????F7??1B", 6 + 2, "NtQueryInformationProcess" },
{ "0F??????????8B????8983????????F6", 9 + 2, "PspCreateProcess" },
{ "8B??8B????8B80BC??????5D", 5 + 2, "PsGetProcessDeubgPort" },
{ "8B??8B????83B8??????????0F", 5 + 2, "PsIsProcessBeingDebugged" },
{ "81????????????74??83BB??????????74", 9 + 2, "NtTerminateProcess" },
{ "85??75??3986????????????56", 4 + 2, "PspTerminateProcess" },
{ "E8????????83BF??????????74??F6", 5 + 2, "PspExitThread" },
{ "89??66??????8DBE????????8B", 6 + 2, "PspProcessDelete" },
{ "64??????????8B????83B8BC????????74", 9 + 2, "ObpCloseHandleTableEntry" },
{ "F6????????????75??398B????????74", 9 + 2, "ObpCloseHandle" },
{ "89????33??399E????????0F", 5 + 2, "MmCreatePeb" },
{ "50E8????????399E????????0F", 6 + 2, "DbgkCreateThread" },
{ "75??8B89????????85??74", 2 + 2, "DbgkExitThread" },
{ "F6????????????75??8B89????????85??74", 9 + 2, "DbgkExitProcess" },
{ "EB??8B89????????3B??74", 2 + 2, "DbgkMapViewOfSection" },
{ "75??83B9??????????74??8B", 2 + 2, "DbgkUnMapViewOfSection" },
{ "89????33??39BE????????0F", 5 + 2, "DbgkpMarkProcessPeb" },
{ "EB??8B????81C3????????39", 5 + 2, "DbgkpCloseObject" },
{ "8B????83A0??????????568B", 3 + 2, "DbgkCopyProcessDebugPort(1)" },
{ "568B????81C6????????83", 4 + 2, "DbgkCopyProcessDebugPort(2)" },
{ "E8????????8B????89B0????????8B", 8 + 2, "DbgkCopyProcessDebugPort(3)" },
{ "568B????81C6????????83", 4 + 2, "DbgkOpenProcessDebugPort" },
{ "FF??????????8B????8B80BC??????89????8B", 9 + 2, "DbgkpQueueMessage" },
{ "FF??8B????05BC??????8B", 5 + 1, "DbgkClearProcessDebugObject" },
{ "C6??????FF??399F????????75", 6 + 2, "DbgkpSetProcessDebugObject(1)" },
{ "8B????8B????8987????????E8", 6 + 2, "DbgkpSetProcessDebugObject(2)" },
{ "85??74??83A7??????????B9", 4 + 2, "DbgkpSetProcessDebugObject(3)" },
{ "C6??????FF??83BF??????????74", 6 + 2, "DbgkpSetProcessDebugObject(4)" },
{ "E8????????EB??83A7??????????C7", 7 + 2, "DbgkpSetProcessDebugObject(5)" },
{ "33??EB??8B81????????32", 4 + 2, "DbgkForwardException" },
{ 0, 0 } };

CODE_SIGN FunctionPointerSign[] =
{
{ "5FC7????????????C7????????????5E", 4, "KiInitializeContextThread->KiThreadStartup" },
{ "5?FF7???FF7???FFB?????????6A0068????????6A00", 16, "PspCreateThread->PspUserThreadStartup" },
{ "84??74??81????????????75??0F", 7, "KiInsertQueueApc->PsExitSpecialApc" },
{ "8B????6A0168????????68????????68????????6A008B??", 16, "PspExitNormalApc->PsExitSpecialApc" },
{ "5?68????????68????????68????????5?FF????5?E8", 12, "PspTerminateThreadByPointer->PsExitSpecialApc" },
{ "8B????6A0168????????68????????68????????6A008B??", 6, "PspExitNormalApc->PspExitNormalApc" },
{ "5?68????????68????????68????????5?FF????5?E8", 2, "PspTerminateThreadByPointer->PspExitNormalApc" },
{ 0, 0 } };

KiDebugPort::KiDebugPort(ULONG oldImageBase, ULONG newImageBase)
{
    mOldImageBase = oldImageBase;
    mNewImageBase = newImageBase;
}

KiDebugPort::~KiDebugPort()
{
    FixInterruptTable(mNewImageBase, mOldImageBase);
}

ULONG KiDebugPort::Char2Hex(CHAR c)
{
    if (c >= '0' && c <= '9')
        return (ULONG) (c - '0');

    if (c >= 'A' && c <= 'Z')
        return (ULONG) (c - 'A' + 10);

    if (c >= 'a' && c <= 'z')
        return (ULONG) (c - 'a' + 10);

    return 0;
}

BOOLEAN KiDebugPort::MatchPattern(PUCHAR Address, PCHAR Pattern)
{
    ULONG Length = strlen(Pattern);
    for (ULONG i = 0; i < Length; i++)
    {
        if (Pattern[i] == '?')
            continue;
        if (((Address[i / 2] >> (4 * ((i + 1) % 2))) & 0x0F) != Char2Hex(Pattern[i]))
            return FALSE;
    }

    return TRUE;
}

BOOLEAN KiDebugPort::ReplaceDebugPortCode(PVOID CodeAddress, ULONG CodeLength)
{
    KiDebug("Search Memory %08X:%08X", CodeAddress, CodeLength);

    if (!MmIsAddressValid(CodeAddress))
    {
        KiDebug("Address Is Invalid!");
        return FALSE;
    }

    for (PUCHAR p = (PUCHAR) CodeAddress; p < (PUCHAR) ((ULONG) CodeAddress + CodeLength); p++)
    {
        for (ULONG i = 0; DebugPortSign[i].Signture; i++)
        {
            if (MatchPattern(p, DebugPortSign[i].Signture))
            {
                KiDebug("Find Signtrue %s %08X->%08X", DebugPortSign[i].Description, p, p + DebugPortSign[i].Offset);

                PMDL mdl = IoAllocateMdl((PVOID) (p + DebugPortSign[i].Offset), 4, FALSE, FALSE, 0);
                MmBuildMdlForNonPagedPool(mdl);
                MmProbeAndLockPages(mdl, KernelMode, IoWriteAccess);
                PLONG newAddress = (PLONG) MmMapLockedPages(mdl, KernelMode);
                InterlockedExchange(newAddress, 0x0000025C);
                MmUnmapLockedPages(newAddress, mdl);
                MmUnlockPages(mdl);
                IoFreeMdl(mdl);
            }
        }
    }
    return TRUE;
}

BOOLEAN KiDebugPort::FixDebugPortCode()
{
    IMAGE_DOS_HEADER* mz_header = (IMAGE_DOS_HEADER*) mNewImageBase;
    IMAGE_NT_HEADERS* pe_header = (IMAGE_NT_HEADERS*) ((ULONG) mz_header + (ULONG) mz_header->e_lfanew);

    ULONG SectionCount = pe_header->FileHeader.NumberOfSections;
    PIMAGE_SECTION_HEADER Section = IMAGE_FIRST_SECTION(pe_header);

    for (ULONG i = 0; i < SectionCount; i++)
    {
        KiDebug("Section Name:%-8s VOffset:%08X VSize:%08X ROffset:%08X RSize:%08X %s", Section[i].Name, Section[i].VirtualAddress, Section[i].Misc.VirtualSize, Section[i].PointerToRawData, Section[i].SizeOfRawData, (Section[i].Characteristics&(IMAGE_SCN_MEM_EXECUTE|IMAGE_SCN_CNT_CODE)) ? "CODE Section" : "");

        if ((Section[i].Characteristics & (IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_CNT_CODE)) && !(Section[i].Characteristics & IMAGE_SCN_MEM_DISCARDABLE))
        {
            ReplaceDebugPortCode((PVOID) (mNewImageBase + (ULONG) Section[i].VirtualAddress), Section[i].Misc.VirtualSize);
        }
    }

    return TRUE;
}

BOOLEAN KiDebugPort::FixCodeReferences()
{
    IMAGE_DOS_HEADER* mz_header = (IMAGE_DOS_HEADER*) mNewImageBase;
    IMAGE_NT_HEADERS* pe_header = (IMAGE_NT_HEADERS*) ((ULONG) mz_header + (ULONG) mz_header->e_lfanew);

    ULONG ImageBase = pe_header->OptionalHeader.ImageBase;

    IMAGE_BASE_RELOCATION* relocate = (IMAGE_BASE_RELOCATION*) ((ULONG) mz_header + (ULONG) pe_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].VirtualAddress);

    if (pe_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC].Size)
    {
        ULONG diff = mNewImageBase - mOldImageBase;
        while (relocate->SizeOfBlock)
        {
            ULONG mempage = (ULONG) mz_header + (ULONG) relocate->VirtualAddress;
            int count = ((ULONG) relocate->SizeOfBlock - (ULONG) sizeof(IMAGE_BASE_RELOCATION)) >> 1;
            USHORT* item = (USHORT*) ((ULONG) relocate + (ULONG) sizeof(IMAGE_BASE_RELOCATION));
            for (int i = 0; i < count; i++)
            {
                int offset = item[i] & 0x0FFF;
                int type = item[i] >> 12;
                if (type == 3)
                {
                    PUCHAR p = (PUCHAR) (mempage + offset);
                    for (ULONG i = 0; FunctionPointerSign[i].Signture; i++)
                    {
                        if (MatchPattern(p - FunctionPointerSign[i].Offset, FunctionPointerSign[i].Signture))
                        {
                            *(PULONG) p += diff;
                            KiDebug("CodeReferences %s %08X->%08X", FunctionPointerSign[i].Description, p, *(PULONG )p);
                        }
                    }
                }
            }
            relocate = (IMAGE_BASE_RELOCATION*) ((ULONG) relocate + (ULONG) relocate->SizeOfBlock);
        }

        return TRUE;
    }

    return FALSE;
}

BOOLEAN KiDebugPort::FixInterruptTable()
{
    return FixInterruptTable(mOldImageBase, mNewImageBase);
}

BOOLEAN KiDebugPort::FixInterruptTable(ULONG mOldImageBase, ULONG mNewImageBase)
{
    SYSTEM_BASIC_INFORMATION SysBaseInfo;

    IMAGE_DOS_HEADER* mz_header = (IMAGE_DOS_HEADER*) mOldImageBase;
    IMAGE_NT_HEADERS* pe_header = (IMAGE_NT_HEADERS*) ((ULONG) mz_header + (ULONG) mz_header->e_lfanew);

    ULONG imageSize = pe_header->OptionalHeader.SizeOfImage;

    if (!NT_SUCCESS(ZwQuerySystemInformation(SystemBasicInformation,&SysBaseInfo,sizeof(SYSTEM_BASIC_INFORMATION),0)))
    {
        return FALSE;
    }

    KiDebug("System Processors Number is %d", SysBaseInfo.NumberProcessors);

    KiDebug("Current Processor %d", KeGetCurrentProcessorNumber());

    typedef KAFFINITY (*pfnKeSetAffinityThread)(PKTHREAD pKThread, KAFFINITY Affinity);
    ULONG OldAffinity = 0;
    for (ULONG nAffinity = 0; nAffinity < SysBaseInfo.NumberProcessors; nAffinity++)
    {
        UNICODE_STRING ustrKeSetAffinityThread;

        pfnKeSetAffinityThread KeSetAffinityThread = NULL;

        RtlInitUnicodeString(&ustrKeSetAffinityThread, L"KeSetAffinityThread");

        KeSetAffinityThread = (pfnKeSetAffinityThread) MmGetSystemRoutineAddress(&ustrKeSetAffinityThread);

        ULONG __affinity = OldAffinity = KeSetAffinityThread(KeGetCurrentThread(), 1 << nAffinity);

        if (!OldAffinity)
            OldAffinity = __affinity;

        if (__affinity)
        {
            KiDebug("CurProcessor %d OldAffinity %d", KeGetCurrentProcessorNumber(), OldAffinity);

            IDTENTRY* idt_entries = GetSystemInterruptTable();
            KiDebug("IDT %08X", idt_entries);

            IDTENTRY* NewIdtEntry = (IDTENTRY*) ExAllocatePool(NonPagedPool, MAX_IDT_ENTRIES * sizeof(IDTENTRY));
            if (!NewIdtEntry)
            {
                KiDebug("Cannot Alloc Pool!");
                return FALSE;
            }

            RtlCopyMemory(NewIdtEntry, idt_entries, MAX_IDT_ENTRIES * sizeof(IDTENTRY));

            ULONG diff = mNewImageBase - mOldImageBase;

            for (int i = 0; i < 255; i++)
            {
                ULONG InterruptProc = MAKELONG(idt_entries[i].LowOffset,idt_entries[i].HiOffset);

                if (InterruptProc >= mOldImageBase && InterruptProc < (mOldImageBase + imageSize))
                {
                    InterruptProc += diff;
                }

                NewIdtEntry[i].LowOffset = (USHORT) InterruptProc;
                NewIdtEntry[i].HiOffset = (USHORT) ((ULONG) InterruptProc >> 16);
            }
            SetSystemInterruptTable(NewIdtEntry);
        }
    }

    UNICODE_STRING ustrKeSetAffinityThread;
    pfnKeSetAffinityThread KeSetAffinityThread = NULL;

    RtlInitUnicodeString(&ustrKeSetAffinityThread, L"KeSetAffinityThread");

    KeSetAffinityThread = (pfnKeSetAffinityThread) MmGetSystemRoutineAddress(&ustrKeSetAffinityThread);
    KeSetAffinityThread(KeGetCurrentThread(), OldAffinity);

    return TRUE;
}

IDTENTRY* KiDebugPort::GetSystemInterruptTable()
{
    IDTINFO idt_info =
    { 0 };
    IDTENTRY* idt_entries = 0;

__asm  sidt  idt_info

    idt_entries = (IDTENTRY*)MAKELONG(idt_info.LowIDTbase,idt_info.HiIDTbase);

        KiDebug("GetIDT %08X %X", idt_entries, idt_info.IDTLimit);

    return idt_entries;
}

BOOLEAN KiDebugPort::SetSystemInterruptTable(IDTENTRY* NewEntry)
{
    IDTINFO idt_info;

    __asm sidt idt_info
    idt_info.LowIDTbase = (USHORT)NewEntry;
    idt_info.HiIDTbase = (USHORT) ((ULONG) NewEntry >> 16);
__asm lidt idt_info

return TRUE;
}
